#include "bytewriter.h"

//
// constructors
//
ByteWriter::ByteWriter(void)
{
	InitVars();
}

ByteWriter::ByteWriter( void* pBuff, DWORD dwSize )
{
	Assign( pBuff, dwSize );
}

//
// destructor
//
ByteWriter::~ByteWriter(void)
{
}

//
// Purpose:
//   Member variable initializer
//
void ByteWriter::InitVars()
{
	pBuffer       = NULL;
	cbBuffer      = 0;
	dwBitsLeft    = dwBitsTotal = 0;
	byBits2Do     = 8;
	byCur         = 0;
	dwIndex       = 0;

	return;
}

//
// Purpose:
//   Class initialization handler
//
BOOL ByteWriter::Assign( void *pBuff, DWORD dwSize )
{
	InitVars();
	if ( !pBuff || !dwSize)
		return FALSE; // ERR

	this->pBuffer  = pBuff;
	this->cbBuffer = dwSize;

	dwBitsLeft = dwBitsTotal = dwSize * 8;

	return TRUE; // OK
}

DWORD ByteWriter::GetUnusedBitCount()
{
	return dwBitsLeft;
}

DWORD ByteWriter::GetTotalBitCount()
{
	return dwBitsTotal;
}

//
// Purpose:
//   Returns the number of accessed bytes
//
DWORD ByteWriter::GetWrittenByteCount()
{
	DWORD  dwcBitsWritten  = (dwBitsTotal - dwBitsLeft);
	DWORD  c = (DWORD)(dwcBitsWritten / 8);

	if ( dwcBitsWritten % 8 ) // round up !
		++c;

	return c;
}

//
// Purpose:
//   MUST be called at the end of the writing operations because it is
//   possible that a unfinished byte must be written to the image
//
void ByteWriter::FinishWriting()
{
	byCur       >>= (byBits2Do-1);
	dwBitsTotal  -= byBits2Do;
	dwBitsLeft   -= byBits2Do;
	byBits2Do     = 8;
	((PBYTE)pBuffer)[ dwIndex++ ] = byCur;
	byCur         = 0;

	return;
}

//
// Purpose:
//   The key member function for the buffer manipulation
//
BOOL ByteWriter::WriteBit( BYTE bit )
{
	if ( dwBitsLeft == 0 )
		return FALSE; // ERR

	// paste bit in current byte
	byCur |= (bit << 7);
	--byBits2Do;
	if ( byBits2Do )
		byCur >>= 1;
	--dwBitsLeft;

	// did we already build a complete byte ? -> paste in buffer
	if ( !byBits2Do )
	{
		((PBYTE)pBuffer)[ dwIndex++ ] = byCur;
		byBits2Do = 8;
		byCur = 0;
	}

	return TRUE; // OK
}

//
// Purpose:
//   Write a 1 byte variable into the buffer
//
BOOL ByteWriter::WriteByte( BYTE byte )
{
	for ( UINT i = 0; i < 8; i++ )
	{
		if ( !WriteBit( byte & 1 ) )
			return FALSE; // ERR
		byte >>= 1;
	}
	return TRUE; // OK
}

//
// Purpose:
//   Write a 2 byte variable into the buffer
//
BOOL ByteWriter::WriteWord( WORD word )
{
	for ( UINT i = 0; i < 16; i++ )
	{
		if ( !WriteBit( (BYTE)(word & 1) ) )
			return FALSE; // ERR
		word >>= 1;
	}
	return TRUE; // OK
}

//
// Purpose:
//   Write a 4 byte variable into the buffer
//
BOOL ByteWriter::WriteDword( DWORD dword )
{
	for ( UINT i = 0; i < 32; i++ )
	{
		if ( !WriteBit( (BYTE)(dword & 1) ) )
			return FALSE; // ERR
		dword >>= 1;
	}
	return TRUE; // OK
}

//
// Purpose:
//   Write a user defined number of bits into the buffer
//
BOOL ByteWriter::WriteBits( DWORD bits, DWORD dwBitCount )
{
	for ( UINT i = 0; i < dwBitCount; i++ )
	{
		if ( !WriteBit( (BYTE)(bits & 1) ) )
			return FALSE; // ERR
		bits >>= 1;
	}
	return TRUE; // OK
}